package tau.dimes.visualization.applet.aviRecord;

import java.awt.Graphics2D;
import java.awt.Image;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;

import javax.imageio.ImageIO;
import javax.jnlp.FileContents;
import javax.jnlp.JNLPRandomAccessFile;

/**
 * MJpeg Generator. Taken from
 * http://www.monceaux.org/coding/MJPEGGenerator.java. Modified to work with
 * jnlp file objects.
 **/
public class MJPEGGenerator
{
	/*
	 * Info needed for MJPEG AVI
	 * 
	 * - size of file minus "RIFF & 4 byte file size"
	 */

	private int width = 0;
	private int height = 0;
	private double framerate = 0;
	private int numFrames = 0;
	private FileContents fileContents = null;
	private JNLPRandomAccessFile aviFile = null;

	private long aviMovieOffset = 0;

	private AVIIndexList indexlist = null;

	/**
	 * Creates a new MJPEGGenerator.
	 * @param file Output file.
	 * @param width Movie width in pixels.
	 * @param height Movie height in pixels.
	 * @param framerate Movie frame rate.
	 * @param numFrames Total number of frames.
	 * @throws IOException
	 */
	public MJPEGGenerator(FileContents file, int width, int height, double framerate, int numFrames)
		throws IOException
	{
		this.fileContents = file;
		this.aviFile = file.getRandomAccessFile("rw");
		this.width = width;
		this.height = height;
		this.framerate = framerate;
		this.numFrames = numFrames;

		RIFFHeader rh = new RIFFHeader();
		writeToAviFile(rh.toBytes());
		writeToAviFile(new AVIMainHeader().toBytes());
		writeToAviFile(new AVIStreamList().toBytes());
		writeToAviFile(new AVIStreamHeader().toBytes());
		writeToAviFile(new AVIStreamFormat().toBytes());
		writeToAviFile(new AVIJunk().toBytes());
		aviMovieOffset = aviFile.getFilePointer();
		writeToAviFile(new AVIMovieList().toBytes());
		indexlist = new AVIIndexList();
	}

	private void writeToAviFile(byte[] bytes) throws IOException
	{
		try
		{
			aviFile.write(bytes);
		}
		catch (IOException e)
		{
			this.fileContents.setMaxLength(this.fileContents.getMaxLength() + bytes.length);
			aviFile.write(bytes);
		}
	}

	/**
	 * Adds a new image to the movie sequence.
	 * @param image An Image instance.
	 * @throws IOException
	 */
	public void addImage(Image image) throws IOException
	{
		byte[] fcc = new byte[] { '0', '0', 'd', 'b' };
		byte[] imagedata = writeImageToBytes(image);
		int useLength = imagedata.length;
		long position = aviFile.getFilePointer();
		int extra = (useLength + (int) position) % 4;
		if (extra > 0)
			useLength = useLength + extra;

		indexlist.addAVIIndex((int) position, useLength);

		writeToAviFile(fcc);
		writeToAviFile(intBytes(swapInt(useLength)));
		writeToAviFile(imagedata);
		if (extra > 0)
		{
			for (int i = 0; i < extra; i++)
				writeToAviFile(0);
		}
		imagedata = null;
	}

	private void writeToAviFile(int i) throws IOException
	{
		try
		{
			aviFile.write(i);
		}
		catch (IOException e)
		{
			fileContents.setMaxLength(fileContents.getMaxLength() + 1024);
			aviFile.write(i);
		}
	}

	/**
	 * Finished the movie and closes the file.
	 * @throws IOException
	 */
	public void finishAVI() throws IOException
	{
		byte[] indexlistBytes = indexlist.toBytes();
		writeToAviFile(indexlistBytes);
		long size = aviFile.length();
		aviFile.seek(4);
		writeToAviFile(intBytes(swapInt((int) size - 8)));
		aviFile.seek(aviMovieOffset + 4);
		writeToAviFile(intBytes(swapInt((int) (size - 8 - aviMovieOffset - indexlistBytes.length))));
		aviFile.close();
	}

	private static int swapInt(int v)
	{
		return (v >>> 24) | (v << 24) | ((v << 8) & 0x00FF0000) | ((v >> 8) & 0x0000FF00);
	}

	private static short swapShort(short v)
	{
		return (short) ((v >>> 8) | (v << 8));
	}

	private static byte[] intBytes(int i)
	{
		byte[] b = new byte[4];
		b[0] = (byte) (i >>> 24);
		b[1] = (byte) ((i >>> 16) & 0x000000FF);
		b[2] = (byte) ((i >>> 8) & 0x000000FF);
		b[3] = (byte) (i & 0x000000FF);

		return b;
	}

	private static byte[] shortBytes(short i)
	{
		byte[] b = new byte[2];
		b[0] = (byte) (i >>> 8);
		b[1] = (byte) (i & 0x000000FF);

		return b;
	}

	private class RIFFHeader
	{
		public byte[] fcc = new byte[] { 'R', 'I', 'F', 'F' };
		public int fileSize = 0;
		public byte[] fcc2 = new byte[] { 'A', 'V', 'I', ' ' };
		public byte[] fcc3 = new byte[] { 'L', 'I', 'S', 'T' };
		public int listSize = 200;
		public byte[] fcc4 = new byte[] { 'h', 'd', 'r', 'l' };

		public RIFFHeader()
		{

		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(fileSize)));
			baos.write(fcc2);
			baos.write(fcc3);
			baos.write(intBytes(swapInt(listSize)));
			baos.write(fcc4);
			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIMainHeader
	{
		/*
		 * 
		 * FOURCC fcc; DWORD cb; DWORD dwMicroSecPerFrame; DWORD
		 * dwMaxBytesPerSec; DWORD dwPaddingGranularity; DWORD dwFlags; DWORD
		 * dwTotalFrames; DWORD dwInitialFrames; DWORD dwStreams; DWORD
		 * dwSuggestedBufferSize; DWORD dwWidth; DWORD dwHeight; DWORD
		 * dwReserved[4];
		 */

		public byte[] fcc = new byte[] { 'a', 'v', 'i', 'h' };
		public int cb = 56;
		public int dwMicroSecPerFrame = 0; // (1 / frames per sec) * 1,000,000
		public int dwMaxBytesPerSec = 10000000;
		public int dwPaddingGranularity = 0;
		public int dwFlags = 65552;
		public int dwTotalFrames = 0; // replace with correct value
		public int dwInitialFrames = 0;
		public int dwStreams = 1;
		public int dwSuggestedBufferSize = 0;
		public int dwWidth = 0; // replace with correct value
		public int dwHeight = 0; // replace with correct value
		public int[] dwReserved = new int[4];

		public AVIMainHeader()
		{
			dwMicroSecPerFrame = (int) ((1.0 / framerate) * 1000000.0);
			dwWidth = width;
			dwHeight = height;
			dwTotalFrames = numFrames;
		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(cb)));
			baos.write(intBytes(swapInt(dwMicroSecPerFrame)));
			baos.write(intBytes(swapInt(dwMaxBytesPerSec)));
			baos.write(intBytes(swapInt(dwPaddingGranularity)));
			baos.write(intBytes(swapInt(dwFlags)));
			baos.write(intBytes(swapInt(dwTotalFrames)));
			baos.write(intBytes(swapInt(dwInitialFrames)));
			baos.write(intBytes(swapInt(dwStreams)));
			baos.write(intBytes(swapInt(dwSuggestedBufferSize)));
			baos.write(intBytes(swapInt(dwWidth)));
			baos.write(intBytes(swapInt(dwHeight)));
			baos.write(intBytes(swapInt(dwReserved[0])));
			baos.write(intBytes(swapInt(dwReserved[1])));
			baos.write(intBytes(swapInt(dwReserved[2])));
			baos.write(intBytes(swapInt(dwReserved[3])));
			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIStreamList
	{
		public byte[] fcc = new byte[] { 'L', 'I', 'S', 'T' };
		public int size = 124;
		public byte[] fcc2 = new byte[] { 's', 't', 'r', 'l' };

		public AVIStreamList()
		{

		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(size)));
			baos.write(fcc2);
			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIStreamHeader
	{
		/*
		 * FOURCC fcc; DWORD cb; FOURCC fccType; FOURCC fccHandler; DWORD
		 * dwFlags; WORD wPriority; WORD wLanguage; DWORD dwInitialFrames; DWORD
		 * dwScale; DWORD dwRate; DWORD dwStart; DWORD dwLength; DWORD
		 * dwSuggestedBufferSize; DWORD dwQuality; DWORD dwSampleSize; struct {
		 * short int left; short int top; short int right; short int bottom; }
		 * rcFrame;
		 */

		public byte[] fcc = new byte[] { 's', 't', 'r', 'h' };
		public int cb = 64;
		public byte[] fccType = new byte[] { 'v', 'i', 'd', 's' };
		public byte[] fccHandler = new byte[] { 'M', 'J', 'P', 'G' };
		public int dwFlags = 0;
		public short wPriority = 0;
		public short wLanguage = 0;
		public int dwInitialFrames = 0;
		public int dwScale = 0; // microseconds per frame
		public int dwRate = 1000000; // dwRate / dwScale = frame rate
		public int dwStart = 0;
		public int dwLength = 0; // num frames
		public int dwSuggestedBufferSize = 0;
		public int dwQuality = -1;
		public int dwSampleSize = 0;
		public int left = 0;
		public int top = 0;
		public int right = 0;
		public int bottom = 0;

		public AVIStreamHeader()
		{
			dwScale = (int) ((1.0 / framerate) * 1000000.0);
			dwLength = numFrames;
		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(cb)));
			baos.write(fccType);
			baos.write(fccHandler);
			baos.write(intBytes(swapInt(dwFlags)));
			baos.write(shortBytes(swapShort(wPriority)));
			baos.write(shortBytes(swapShort(wLanguage)));
			baos.write(intBytes(swapInt(dwInitialFrames)));
			baos.write(intBytes(swapInt(dwScale)));
			baos.write(intBytes(swapInt(dwRate)));
			baos.write(intBytes(swapInt(dwStart)));
			baos.write(intBytes(swapInt(dwLength)));
			baos.write(intBytes(swapInt(dwSuggestedBufferSize)));
			baos.write(intBytes(swapInt(dwQuality)));
			baos.write(intBytes(swapInt(dwSampleSize)));
			baos.write(intBytes(swapInt(left)));
			baos.write(intBytes(swapInt(top)));
			baos.write(intBytes(swapInt(right)));
			baos.write(intBytes(swapInt(bottom)));
			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIStreamFormat
	{
		/*
		 * FOURCC fcc; DWORD cb; DWORD biSize; LONG biWidth; LONG biHeight; WORD
		 * biPlanes; WORD biBitCount; DWORD biCompression; DWORD biSizeImage;
		 * LONG biXPelsPerMeter; LONG biYPelsPerMeter; DWORD biClrUsed; DWORD
		 * biClrImportant;
		 */

		public byte[] fcc = new byte[] { 's', 't', 'r', 'f' };
		public int cb = 40;
		public int biSize = 40; // same as cb
		public int biWidth = 0;
		public int biHeight = 0;
		public short biPlanes = 1;
		public short biBitCount = 24;
		public byte[] biCompression = new byte[] { 'M', 'J', 'P', 'G' };
		public int biSizeImage = 0; // width x height in pixels
		public int biXPelsPerMeter = 0;
		public int biYPelsPerMeter = 0;
		public int biClrUsed = 0;
		public int biClrImportant = 0;

		public AVIStreamFormat()
		{
			biWidth = width;
			biHeight = height;
			biSizeImage = width * height;
		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(cb)));
			baos.write(intBytes(swapInt(biSize)));
			baos.write(intBytes(swapInt(biWidth)));
			baos.write(intBytes(swapInt(biHeight)));
			baos.write(shortBytes(swapShort(biPlanes)));
			baos.write(shortBytes(swapShort(biBitCount)));
			baos.write(biCompression);
			baos.write(intBytes(swapInt(biSizeImage)));
			baos.write(intBytes(swapInt(biXPelsPerMeter)));
			baos.write(intBytes(swapInt(biYPelsPerMeter)));
			baos.write(intBytes(swapInt(biClrUsed)));
			baos.write(intBytes(swapInt(biClrImportant)));
			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIMovieList
	{
		public byte[] fcc = new byte[] { 'L', 'I', 'S', 'T' };
		public int listSize = 0;
		public byte[] fcc2 = new byte[] { 'm', 'o', 'v', 'i' };

		// 00db size jpg image data ...

		public AVIMovieList()
		{

		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(listSize)));
			baos.write(fcc2);
			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIIndexList
	{
		public byte[] fcc = new byte[] { 'i', 'd', 'x', '1' };
		public int cb = 0;
		public ArrayList<AVIIndex> ind = new ArrayList<AVIIndex>();

		public AVIIndexList()
		{

		}

		public void addAVIIndex(int dwOffset, int dwSize)
		{
			ind.add(new AVIIndex(dwOffset, dwSize));
		}

		public byte[] toBytes() throws IOException
		{
			cb = 16 * ind.size();

			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(cb)));
			for (int i = 0; i < ind.size(); i++)
			{
				AVIIndex in = (AVIIndex) ind.get(i);
				baos.write(in.toBytes());
			}

			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIIndex
	{
		public byte[] fcc = new byte[] { '0', '0', 'd', 'b' };
		public int dwFlags = 16;
		public int dwOffset = 0;
		public int dwSize = 0;

		public AVIIndex(int dwOffset, int dwSize)
		{
			this.dwOffset = dwOffset;
			this.dwSize = dwSize;
		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(dwFlags)));
			baos.write(intBytes(swapInt(dwOffset)));
			baos.write(intBytes(swapInt(dwSize)));
			baos.close();

			return baos.toByteArray();
		}
	}

	private class AVIJunk
	{
		public byte[] fcc = new byte[] { 'J', 'U', 'N', 'K' };
		public int size = 1808;
		public byte[] data = new byte[size];

		public AVIJunk()
		{
			Arrays.fill(data, (byte) 0);
		}

		public byte[] toBytes() throws IOException
		{
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(fcc);
			baos.write(intBytes(swapInt(size)));
			baos.write(data);
			baos.close();

			return baos.toByteArray();
		}
	}

	private byte[] writeImageToBytes(Image image) throws IOException
	{
		BufferedImage bi = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		Graphics2D g = bi.createGraphics();
		g.drawImage(image, 0, 0, width, height, null);
		ImageIO.write(bi, "jpg", baos);
		baos.close();
		bi = null;
		g = null;
		
//		Iterator<ImageWriter> iter = ImageIO.getImageWritersByFormatName("jpeg");
//		ImageWriter writer = iter.next();
//		ImageWriteParam iwp = writer.getDefaultWriteParam();
//		iwp.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
//		//iwp.setCompressionType("BI_JPEG");
//		iwp.setCompressionQuality(1);
//		ByteArrayOutputStream baos = new ByteArrayOutputStream();
//		IIOImage iioimage = new IIOImage(bi, null, null);
//		writer.setOutput(new MemoryCacheImageOutputStream(baos));
//		writer.write(null, iioimage, iwp);
//		writer.dispose();		

		return baos.toByteArray();
	}
}
